
#include "glodef.h"
#include "errlog.h"
#include "vecs.h"
#include "grid3d.h"
#include "fe.h"
#include "f3io.h"

#include "glovar.h"
#include "soah_common.h"

#include <stdlib.h>
#include <stdio.h>

int MAX_PERIODS=10000;

static void backup_state_dir(void);
static void backup_adjoint_dir(void);
/* static */ void erase_backup(const char *dir);

static void log_period_norms(double norm_u, double norm_p, int period);
static void log_adjoint_period_norms(double norm_u, double norm_p, int period);

int __solve_period(vector u_init, scalar p_init, int period,
    double *norm_u, double *norm_p)__attribute__((warn_unused_result));
int __solve_adjoint(vector u_init, scalar p_init, int period, vector v_finl,
    scalar q_finl, vector h, double *norm_h)__attribute__((warn_unused_result));
int __solve_shape_adjoint_period(vector v_finl, scalar q_finl, int period,
    double *norm_v, double *norm_q)__attribute__((warn_unused_result));

/*
 #define _ADJOINT_B11
 */

int solve_state_periodic(vector u_init, scalar p_init)
{
    int ret=TRUE;
    int period;
    double norm_u, norm_p;

    PrintLog("Solving periodic state problem by marching.\n");

#ifdef  _ADJOINT_B11
    vector v_finl, h;
    scalar q_finl;

    v_finl=alloc_vec(get_N_vec(u));
    h=alloc_vec(get_N_vec(u));
    q_finl=alloc_scal(get_N_scal(p));
#endif

    for(period=1; period<=MAX_PERIODS; period++)
    {

        ret=__solve_period(u_init, p_init, period, &norm_u, &norm_p);
        if(!ret)
            break;
#ifdef  _ADJOINT_B11
        __solve_adjoint(u_init, p_init, period, v_finl, q_finl, h, NULL);
        /* prepare initial condition for next period */
        curr_ti=total_ti;
        load_state_time(state_dir());
#endif
        /*
         compute_shape_gradient();
         print_vec(sg_z, "Z.txt");
         print_scal(sg_xi, "xi.txt");
         */

        if(norm_u<PERIODIC_TOL)
            break;

        copy_vec(u, u_init);
        copy_scal(p, p_init);
    }

    erase_backup(state_dir());

#ifdef  _ADJOINT_B11
    free_vec(v_finl);
    free_vec(h);
    free_scal(q_finl);
#endif
    return ret;
}

int solve_shape_adjoint_periodic(vector v_finl, scalar q_finl)
{
    int ret=TRUE;
    int period;
    double norm_v, norm_q;

    PrintLog("Solving periodic adjoint problem by marching.\n");


    for(period=1; period<=MAX_PERIODS; period++)
    {

        ret=__solve_shape_adjoint_period(v_finl, q_finl, period,
            &norm_v, &norm_q);
        if(!ret)
            break;

        if(norm_v<PERIODIC_TOL)
            break;

        /* prepare final condition for next period */
        curr_ti=0;
        load_adjoint_time(adjoint_dir());
        copy_vec(v, v_finl);
        copy_scal(q, q_finl);
    }

    erase_backup(adjoint_dir());
    return ret;
}

int solve_state_periodic_adjoint(vector u_init, scalar p_init)
{
    int ret=TRUE;
    int period, uN=get_N_vec(u), pN=get_N_scal(p);
    double norm_u, norm_p;

    vector v_finl=alloc_vec(uN);
    vector h=alloc_vec(uN);
    scalar q_finl=alloc_scal(pN);

    for(period=1; period<=MAX_PERIODS; period++)
    {

        ret=__solve_period(u_init, p_init, period, &norm_u, &norm_p);
        if(!ret)
            break;
        /*
         * if(fmax(norm_u, norm_p) < PERIODICL_TOL)
         * break;
         */
        if(norm_u<PERIODIC_TOL)
            break;

        /* prepare initial condition for next period */
        if(flags.need_adjoint)
        {
            ret=__solve_adjoint(u_init, p_init, period, v_finl, q_finl, h, NULL);
            if(!ret)
                break;

            // for(int i=0;i<3*u->N;i++) u_init->X[i] -= h->X[i];
            axpy_vec(u_init, -1.0, h);
        }
        else /* do nothing */
        {
            copy_vec(u, u_init);
            copy_scal(p, p_init);
            /* if (period > 2)
             * {
             * int i;
             * for(i=0;i<3*u->N;i++)
             * u->X[i] = 3./2.*u->X[i]- 1./2.*u_init->X[i];
             * for(i=0;i<p->N;i++)
             * p->F[i] = 3./2.*p->F[i]- 1./2.*p_init->F[i];
             * }*/
        }

        clean_scal(phi);
    }

    free_vec(v_finl);
    free_vec(h);
    free_scal(q_finl);
    return ret;
}

static vector u_init_c;

int solve_state_periodic_armijo(vector u_init, scalar p_init)
{
    int ret=TRUE;
    int period, step, uN=get_N_vec(u), pN=get_N_scal(p);
    double Fc, norm_u, norm_p, norm_h, norm_h0;

    double ta=PERIODIC_TOL,
        tr=ta*1e2;
    double alpha=1e-3;
    double beta=0.8;
    double lambda=0.5;

    vector v_finl=alloc_vec(uN);
    vector h=alloc_vec(uN);
    scalar q_finl=alloc_scal(pN);

    PrintLog("Solving periodic state problem by armijo.\n");

    Fc=1e99;
    norm_h0=0;
    norm_h=0;

    u_init_c=alloc_vec(uN);
    copy_vec(u_init, u_init_c);

    for(step=period=1; period<=MAX_PERIODS; period++)
    {
        double F;

        ret=__solve_period(u_init, p_init, period, &norm_u, &norm_p);
        if(!ret)
            break;
        F=0.5*sqr(norm_u);

        if(F-Fc < -alpha*lambda*norm_h)
        {
            /* accept step */
            lambda=0.5;
            Fc=F;
            copy_vec(u_init, u_init_c);
            ret=__solve_adjoint(u_init, p_init, step, v_finl, q_finl, h, &norm_h);
            if(!ret)
                break;
            if(step==1)
                norm_h0=norm_h;
            if(norm_h<tr*norm_h0+ta)
                break;
            step++;
        }
        else
        {
            /* reject step */
            lambda*=beta;
        }
        // for(int i=0;i<3*u->N;i++) u_init->X[i] = u_init_c->X[i] - lambda * h->X[i];
        waxpy_vec(u_init, -lambda, h, u_init_c);
    }

    free_vec(u_init_c);
    free_vec(v_finl);
    free_vec(h);
    free_scal(q_finl);
    return ret;
}







/****************************************************************************/
/*                                                                          */
/*                       =========================                          */
/*                                                                          */
/****************************************************************************/

static void log_period_norms(double norm_u, double norm_p, int period)
{
    PrintLogNorm("u_i-u_f", "L2", norm_u);
    PrintLogNorm("u_f", "L2", current_solution_norm[NORM_ERR_U_L2]);
    PrintLogNorm("p_i-p_f", "L2", norm_p);
    PrintLogNorm("p_f", "L2", current_solution_norm[NORM_ERR_P_L2]);
    PrintLog("       %10s   %-20s = %.8g\n", "total_cost", "", total_cost);
    PRINTF("period  %3d : ", period);
    PRINTF(" ||u_i-u_f|| = %10.4e", norm_u);
    PRINTF(" ||p_i-p_f|| = %10.4e", norm_p);
    PRINTF("\n");
}

static void log_adjoint_period_norms(double norm_v, double norm_q,
    int period)
{
    PrintLogNorm("v_i-v_f", "L2", norm_v);
    PrintLogNorm("q_i-q_f", "L2", norm_q);
    PRINTF("adjoint %3d : ", period);
    PRINTF(" ||v_i-v_f|| = %10.4e", norm_v);
    PRINTF(" ||q_i-q_f|| = %10.4e", norm_q);
    PRINTF("\n");
}

static int ignor;

void erase_backup(const char *dir)
{
    if(mpi_my_id!=0)
        return;

    char cmd[FILENAME_MAX+10];
    sprintf(cmd, "rm -r -f %s_backup", dir);
    ignor=system(cmd);
}

static void backup_dir(const char *dir)
{
    if(mpi_my_id!=0)
        return;

    char back_dir[FILENAME_MAX];
    char cmd[2*FILENAME_MAX+10];
    sprintf(back_dir, "%s_backup", dir);
    sprintf(cmd, "rm -r -f %s && [ -d %s ] && mv -f %s %s",
        back_dir, dir, dir, back_dir);
    ignor=system(cmd);
}

static void backup_adjoint_dir(void)
{
    backup_dir(adjoint_dir());
}

static void backup_state_dir(void)
{
    backup_dir(state_dir());
}

/****************************************************************************/
/*                                                                          */
/*                       =========================                          */
/*                                                                          */

/****************************************************************************/

int __solve_period(vector u_init, scalar p_init, int period,
    double *norm_u, double *norm_p)
{
    static char state_log[FILENAME_MAX];
    static char curnt_log[FILENAME_MAX];
    if(dump_bin>0)
        sprintf(state_log, "%s/state_period_%03d.log",
        logs_dir(), period);
    else
        sprintf(state_log, "state_period_%03d.log", period);
    PrintLog("Logging solution of state period %d to %s\n",
        period, state_log);

    GetLogfilePath(curnt_log, FILENAME_MAX);
    InitLogfile(state_log);

    backup_state_dir();

    clean_scal(phi);
    int ret=solve_state_ivp(u_init, p_init);
    PRINTF("\r%60s\r", " ");

    AppendLogfile(curnt_log);

    (*norm_u)=norm_vec_diff_l2(grid, elemU, u, u_init);
    (*norm_p)=norm_scal_diff_l2(grid, elemP, p, p_init);
    log_period_norms(*norm_u, *norm_p, period);
    return ret;
}

/****************************************************************************/
/*                                                                          */
/*                       =========================                          */
/*                                                                          */

/****************************************************************************/

int __solve_shape_adjoint_period(vector v_finl, scalar q_finl,
    int period, double *norm_v, double *norm_q)
{
    static char adjoint_log[FILENAME_MAX];
    static char current_log[FILENAME_MAX];
    if(dump_bin>0)
        sprintf(adjoint_log, "%s/adjoint_period_%03d.log",
        logs_dir(), period);
    else
        sprintf(adjoint_log, "adjoint_period_%03d.log", period);
    PrintLog("Logging solution of adjoint period %d to %s\n",
        period, adjoint_log);
    GetLogfilePath(current_log, FILENAME_MAX);
    InitLogfile(adjoint_log);

    backup_adjoint_dir();

    clean_scal(qhi);
    int ret=solve_adjoint_shape_fvp(v_finl, q_finl);
    PRINTF("\r%60s\r", " ");

    AppendLogfile(current_log);

    (*norm_v)=norm_vec_diff_l2(grid, elemU, v, v_finl);
    (*norm_q)=norm_scal_diff_l2(grid, elemP, q, q_finl);
    log_adjoint_period_norms(*norm_v, *norm_q, period);
    return ret;
}

/****************************************************************************/
/*                       =========================                          */

/****************************************************************************/


int __solve_adjoint(vector u_init, scalar p_init, int period,
    vector v_finl, scalar q_finl, vector h, double *norm_h)
{
    static char adjoint_log[FILENAME_MAX];
    static char current_log[FILENAME_MAX];
    if(dump_bin>0)
        sprintf(adjoint_log, "%s/adjoint_period_%03d.log",
        logs_dir(), period);
    else
        sprintf(adjoint_log, "adjoint_period_%03d.log", period);
    PrintLog("Logging solution of adjoint period %d to %s\n",
        period, adjoint_log);
    GetLogfilePath(current_log, FILENAME_MAX);
    InitLogfile(adjoint_log);

    // for(int i=0;i<3*u->N;i++) v_finl->X[i] = u_init->X[i] - u->X[i];
    waxpy_vec(v_finl, -1.0, u, u_init);
#if defined(CS_INITIAL_PRESSURE)
    clean_scal(q_finl);
#endif
    (void)p_init; /* avoid compiler warning */

    int ret=solve_adjoint_periodic_fvp(v_finl, q_finl);
    PRINTF("\r%60s\r", " ");

    /* this is the gradient */
    // for(int i=0;i<3*u->N;i++) h->X[i] = v_finl->X[i] - v->X[i];
    waxpy_vec(h, -1.0, v, v_finl);

    AppendLogfile(current_log);

    {
        double norm=norm_vec_diff_l2(grid, elemU, h, NULL);
        PrintLogNorm(" h ", "L2", norm);
        PRINTF("adjoint %3d :  ||   h   || = %10.4e\n", period, norm);
        if(norm_h) *norm_h=norm;
    }
    return ret;
}

